Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request ensures that Arrow queries returning zero rows still yield a single empty batch so that callers can retrieve the schema. It introduces a helper function _empty_batch and updates download_arrow_from_job to track whether any batches have been yielded, along with adding corresponding unit and system tests. A review comment suggests guarding the empty batch yield when reader is None with if not yielded to prevent appending an extra empty batch to a non-empty result set.
| if reader is None: | ||
| # An empty destination table has no streams. | ||
| yield _empty_batch(pa_schema, row_iterator._schema) | ||
| return |
There was a problem hiding this comment.
If reader is None but yielded is already True (e.g., some batches were already yielded in Step 1), yielding an empty batch here would append an extra empty batch to a non-empty result set. We should guard this with if not yielded: to ensure we only yield an empty batch when no rows have been yielded at all.
| if reader is None: | |
| # An empty destination table has no streams. | |
| yield _empty_batch(pa_schema, row_iterator._schema) | |
| return | |
| if reader is None: | |
| # An empty destination table has no streams. | |
| if not yielded: | |
| yield _empty_batch(pa_schema, row_iterator._schema) | |
| return |
There was a problem hiding this comment.
reader is only None on the explicit-destination path, and explicit destinations always go through jobs.insert, where we wait with getQueryResults(maxResults=0) so there's no inline Arrow batch. yielded is always False there.
825fed7 to
b501684
Compare
05834ef to
d77849b
Compare
d77849b to
dbc30a0
Compare
0d21c36 to
4b1ddd1
Compare
a7b0ad3 to
263c70f
Compare
263c70f to
d7a6b4c
Compare
d7a6b4c to
32b3c10
Compare
32b3c10 to
f277a64
Compare
| # this package's table module. | ||
| from google.cloud.bigquery import _pandas_helpers | ||
|
|
||
| pa_schema = _pandas_helpers.bq_to_arrow_schema(bq_schema) |
There was a problem hiding this comment.
Is this necessary? Do we not get an arrow schema in the jobs.query response if the query completes with empty results.
There was a problem hiding this comment.
I filed b/573768143. Until this is addressed, we'll need to ensure we return an empty table when the response is empty.
I believe the empty schema case (this function) is not relevant to jobs.query but to the case where ReadRows is called after jobs.insert.
Correction: you're right that jobs.query returns arrowSchema when the query completes, even with 0 rows, so this isn't used there. It's for empty results read from the job's _default stream: ReadRows sends no messages at all for 0 rows, so there's no schema. That happens after jobs.insert with an anonymous destination, and also after jobs.query when the query doesn't finish in time and we wait. Checked all of these against BigQuery. Explicit destinations get the schema from the read session.
A query that returned no rows with ARROW yielded no batches, so to_arrow() fell back to the iterator's schema, which jobs.query omits for ARROW, and returned a table without columns. jobs.query sends arrowSchema even for zero rows, but the default job stream sends nothing for them. to_arrow_iterable() now yields one empty RecordBatch if no rows arrive, like _EmptyRowIterator already does. It has the Arrow schema from the first page or the read session; otherwise the schema from jobs.getQueryResults, which the jobs.insert path has; otherwise no columns. to_arrow() and pyarrow.Table.from_batches() keep that schema.
f277a64 to
450fe2f
Compare
ohmayr
added
the
kokoro:force-run
yoshi-kokoro
removed
the
kokoro:force-run
This branch has not been deployed
3 participants
Yields one empty
RecordBatchwith the query's schema when Arrow results have no rows, soto_arrow()no longer returns a table without columns.jobs.getQueryResults, else no columns, like_EmptyRowIterator.